from sklearn.decomposition import PCA
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
import pandas as pd
import config

def get_data(url):
    return pd.read_csv(url)

def do_cluster(data_df):
    num_data_arr = data_df.values.astype('float')
    y_pred = KMeans(max_iter=20, n_clusters=3,init='k-means++', tol=1e-12).fit_predict(num_data_arr)
    return y_pred

def do_PCA(data_df):
    return PCA(n_components=2).fit_transform(data_df)

def get_data_clustered(data_dict, y_pred):
    temp = {}
    for i in range(len(y_pred)):
        temp[i] = y_pred[i]
    data_dict['y'] = temp
    return pd.DataFrame(data_dict)

def main(config):
    for _ in config['repository_list']:
        data = {}
        data_df = get_data(config['data_path']['final'].format(_.split('/')[-1])+'.csv')
        num_data_df = data_df[['commit', 'issue', 'pullrequest', 'issuecomment', 'commitcomment', 'contri']]
        X_pca = do_PCA(num_data_df)
        y_pred = do_cluster(num_data_df)
        # 可视化聚类结果（PCA降维数据视图按照聚类结果上色）
        plt.scatter(X_pca[:, 0], X_pca[:, 1], c=y_pred,label="PCA")
        plt.savefig(config['data_path']['cluster'].format(_.split('/')[-1])+'.png')
        # 保存结果
        data_clustered_df = get_data_clustered(data_df.to_dict(), y_pred)
        data_clustered_df.to_csv(config['data_path']['cluster'].format(_.split('/')[-1])+'.csv')

if __name__ == '__main__':
    _config = config.get_config()
    main(_config)